from typing import List
import pathlib
import os

from pymodels.utils.BookReader import BasicBookReader
from pymodels.models.dpr import DPRVectorizer, tinybertVectorizer
from pymodels.utils.Indexer import FaissIndexerHelper
from pluginUtils import BasicReaderPlugin
from utils.idle import idle

import fitz
from paddleocr import PaddleOCR, PPStructure
import cv2
import tqdm
from transformers import BertConfig


class ReaderPlugin(BasicReaderPlugin, BasicBookReader):
    def __init__(self, modelname: str = "zerohell/tinydpr-acc_0.315-bs_307", device: str = "cuda:0", size: int = 512):
        """
        设置PDF阅读器。
        :param modelname: DPR模型路经。
        :param device: 推理时所使用的处理器。'cpu'或'cuda:n'。推荐采用：cuda:0
        :param size: 读取文本时的每页的大小。越小，推理器的结果可能越精确。推荐值：512.
        :return:
        """
        self.modelname = modelname
        self.device = device
        dim = BertConfig.from_pretrained(modelname).hidden_size
        indexer = FaissIndexerHelper(dim=dim)
        self.page_size = 0
        self.size = size

        BasicBookReader.__init__(self, indexer=indexer, booktype='.pdf', cache=True, arduino=True)
        BasicReaderPlugin.__init__(self)

    def hibernate(self):
        BasicBookReader.del_vectorizer(self)

    def wake(self):
        dpr = tinybertVectorizer(self.modelname)
        dpr.model.to(self.device)
        BasicBookReader.set_vectorizer(self, dpr)

    def file2pages(self, filename: str) -> List[str]:  # noqa
        """
        :param filename: 要处理的文件名称。
        :return: 页形成的数组。会包括一个空页，用于判断是否为不相关的问题。
        """
        size = self.size
        blocks = ['']
        overlap = size // 2
        self.page_size = size
        filename = pathlib.Path(filename)
        self.readbook_name = filename.name
        content = self._pdf2txt(filename)
        blocks = [content[i: i + size] for i in range(0, len(content), overlap)]
        return blocks

    def _pdf2txt(self, filename: pathlib.Path):
        if getattr(self, "ocr", None) is None:
            self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=True)
        if getattr(self, "table_engine", None) is None:
            self.table_engine = PPStructure(table=False, ocr=False)
        content = ""
        cachefile_dir = filename.parent.joinpath(filename.stem + '.cache')
        if not cachefile_dir.is_dir():
            cachefile_dir.mkdir()
        if cachefile_dir.joinpath("content").is_file():
            with open(cachefile_dir.joinpath("content"), mode='r', encoding='utf8') as file:
                content = file.read()
            return content

        doc = fitz.open(filename)

        def getcroresult(path):
            result_ = self.ocr.ocr(path, cls=True)
            texts = [res[1][0] for res in result_[0]]
            text_ = "\n".join(texts)
            return text_

        temp1 = "1.jpg"
        temp2 = "2.jpg"

        for i in tqdm.tqdm(range(len(doc))):
            page = doc.load_page(i)
            pix = page.get_pixmap()
            pix.save(temp1)
            img = cv2.imread(temp1)
            result = self.table_engine(img)
            for line in result:
                if line['type'] == 'text':
                    text_img = line['img']
                    cv2.imwrite(temp2, text_img)
                    text = getcroresult(temp2)
                    content += text
                # else:
                #     print("abandon", line['type'])
        os.remove(temp1)
        os.remove(temp2)
        with open(cachefile_dir.joinpath("content"), mode='w', encoding='utf8') as file:
            file.write(content)
        return content

    def get_page_size(self) -> int:
        return self.page_size

    def search(self, query: str, k: int):
        return BasicBookReader.search(self, query, k=k, return_page_content=True, return_cos_similarity=True)

    def readbook(self, bookpath: str):
        return BasicBookReader.readbook(self, filename=bookpath)

    def get_readfile_name(self):
        return BasicBookReader.get_readfile_name(self)

    def get_introduction(self):
        return """
        作者：xzh
        日期：2023年4月24日
        邮箱：kodderopert@163.com
        """


if __name__ == "__main__":
    plugin = ReaderPlugin(**{
            "dprmodel": "D:\\static\\trained model\\dpr-bert-acc@0.961\\model.param.pt",
            "tokenpath": "D:\\static\\trained model\\mengzi"
    })
    plugin.readbook("bookshelf/转换测试.pdf")
    idle(globals())
